import os
import sys
import rand
import datetime
import torch
import numpy as np
import hydra
from omegaconf import OmegaConf

from algorithms import agents
from utils.utils import Logger
from utils.wrappers import GrayScaleObservationWrapper, ResizeObservationWrapper, FrameStackWrapper, \
    PettingzooToGymnasiumEnv, create_pettingzoo_env


def make_save(cfg):
    time = str(datetime.datetime.now().replace(microsecond=0).strftime("%Y.%m.%d.%H.%M.%S"))
    cfg.save_path = os.path.join(cfg.base_path, time)
    if not os.path.exists(cfg.save_path):
        os.makedirs(cfg.save_path)
    if cfg.model:
        cfg.model_path = os.path.join(cfg.save_path, "models")
        if not os.path.exists(cfg.model_path):
            os.makedirs(cfg.model_path)


def make_env(env_name, gray_scale=True, shape=84, num_stack=4):
    env = create_pettingzoo_env(env_name)
    env = ResizeObservationWrapper(env, shape=shape)
    if gray_scale:
        env = GrayScaleObservationWrapper(env)
    env = FrameStackWrapper(env, num_stack=num_stack)
    # env = PettingzooToGymnasiumEnv(env)
    return env


def run(args, stdout):
    make_save(args)
    OmegaConf.save(args, os.path.join(args.save_path, "config.yaml"))
    sys.stdout = Logger(stdout, os.path.join(args.save_path, "logs.txt"))
    print("============================================================")
    print("saving at:", args.save_path)
    # create train env and eval env
    env = make_env(args.env, args.gray_scale, args.shape, args.num_stack)
    eval_env = make_env(args.env, args.gray_scale, args.shape, args.num_stack)

    device = torch.device("cuda" if torch.cuda.is_available() and args.cuda else "cpu")

    # create agent
    agent = agents[args.algorithms.algo_name](args, env, eval_env, device)
    agent.run()
    print("============================================================")
    print("saving at:", args.save_path)
    print("============================================================")
    sys.stdout.close()


@hydra.main(config_path='configs', config_name='config', version_base=None)
def main(cfg):
    stdout = sys.stdout
    for r in range(cfg.run):
        run(cfg, stdout)


if __name__ == "__main__":
    main()
